/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.example.grid;

import java.util.List;

import net.javlov.Action;
import net.javlov.Agent;
import net.javlov.Environment;
import net.javlov.State;
import net.javlov.VectorState;

public class SimpleGridWorld implements Environment<double[]> {

	/**
	 * Current state the agent is in
	 */
	protected VectorState currentState;
	
	/**
	 * Whether the world wraps around
	 */
	protected boolean wrap;
	
	/**
	 * Width and height in number of cells
	 */
	protected int width, height;
	
	/**
	 * Goal state, if any
	 */
	protected VectorState goalState;
	
	/**
	 * List of user-defined objects
	 */
	protected List<?> objects;
	
	/**
	 * 
	 * @param width
	 * @param height
	 * @param wrap
	 */
	public SimpleGridWorld(int width, int height, boolean wrap) {
		this.width = width;
		this.height = height;
		this.wrap = wrap;
		currentState = new VectorState(null);
	}
	
	public VectorState getCurrentState() {
		return currentState;
	}

	public boolean isWrap() {
		return wrap;
	}

	public int getWidth() {
		return width;
	}

	public int getHeight() {
		return height;
	}

	public VectorState getGoalState() {
		return goalState;
	}

	@Override
	public double executeAction(Action<double[]> act, Agent<double[]> a) {
		act.execute(a);
		if ( goalReached() )
			currentState.setTerminal(true);
		return -1;
	}

	@Override
	public State<double[]> getObservation(Agent<double[]> a) {
		return currentState;
	}

	@Override
	public int getObservationDim() {
		return 2;
	}

	@Override
	public void init() {
		reset();		
	}

	@Override
	public void reset() {	
		currentState = new VectorState( new double[]{
				Math.floor(Math.random()*width),
				Math.floor(Math.random()*height)
		} );
	}
	
	public boolean goalReached() {
		return currentState.equals(goalState);
	}
	
	public void setGoal(int x, int y) {
		goalState = new VectorState(new double[]{x,y});
	}
	
	public boolean move(Direction d) {
		double[] p = currentState.getData(),
				pnew = d.go(p);
		boolean m = checkWrap(p,pnew);
		currentState = new VectorState(pnew);
		return m;
	}
	
	protected boolean checkWrap(double[] p, double[] pnew) {
		boolean m = true;
		if ( pnew[0] < 0 || pnew[1] < 0 || pnew[0] > width-1 || pnew[1] > height-1 ) {
			if ( wrap ) {
				//move to other side of grid
				pnew[0] = ( pnew[0] < 0 ? width-1 : pnew[0] % width );
				pnew[1] = ( pnew[1] < 0 ? height-1 : pnew[1] % height );
			} else {
				//stay in the same place
				m = false;
				pnew[0] = p[0];
				pnew[1] = p[1];
			}
		}
		return m;
	}
	
	public void setObjects(List<?> objs) {
		objects = objs;
	}
	
	public List<?> getObjects() {
		return objects;
	}
}
